{-# LANGUAGE ParallelListComp #-}
{-# LANGUAGE PackageImports #-}

module Translator(translate) where

import qualified Data.Map as Map
import "mtl" Control.Monad.State.Strict(get, put, gets, modify)
import Control.Monad(forM, forM_)
import Types
import Env
import Debug.Trace
import Printer(printExpr, printStmt)
import Misc


x # f = f x
infixl 1 #

debug :: (Monad m) => String -> m ()
debug s = trace s (return ())


-- translate expression
transExpr :: Expr -> Trans (Expr, Type)
transExpr (V x) = do
    env <- get
    case env # lenv # Map.lookup x of
        Nothing -> error ("trying to use unassigned variable: " ++ x)
        Just t -> return (V x, t)

transExpr (C x) = return (C x, Float)
transExpr (Ci x) = return (Ci x, Int)
transExpr (S x) = return (S x, String)
transExpr T = return (T, Boolean)
transExpr F = return (F, Boolean)

transExpr e@(Func name args) = do
    env <- get
    case env # fenv # Map.lookup name of
        Nothing -> case env # uenv # Map.lookup name of
            Nothing -> error $ "undefined function: " ++ name
            Just _ -> transUserCall name args
        Just t -> transFunCall name t args

transExpr (Op name e1 e2) = do
    (e1', t1) <- transExpr e1
    (e2', t2) <- transExpr e2
    let t = join t1 t2
    if isNumeric t
        then return (Op name (cast t1 t e1') (cast t2 t e2'), t)
        else error $ "bad operator arguments in " ++ show (Op name e1 e2)
        
transExpr (Cast t e) = do
    (e', t') <- transExpr e 
    return (cast t' t e', t)

transExpr (Dot expr x) =
    let
        f 2 = Float2
        f 3 = Float3
        f 4 = Float4

        g 2 = Int2
        g 3 = Int3
        g 4 = Int4

        trans e t | t `elem` [Float2, Float3, Float4] =
            case length x of
                1 -> (Dot e x, Float) -- vector element access
                n -> (Func ("get_" ++ x) [e], f n) -- swizzle operator
    
        trans e t | t `elem` [Int2, Int3, Int4] =
            case length x of
                1 -> (Dot e x, Int) -- vector element access
                n -> (Func ("iget_" ++ x) [e], g n) -- swizzle operator

        trans e Dataset = case e of
            V dataset -> (V (dataset ++ "_size"), Float3)
            _ -> error "compiler badness"
    
        trans _ _ = error $ "dot applied to wrong type: " ++ printExpr expr
    in do
        (e', t') <- transExpr expr
        return (trans e' t')

-- helper function to translate multiple expressions
transExprs :: [Expr] -> Trans ([Expr], [Type])
transExprs es = do
    es' <- forM es transExpr
    return (unzip es')


-- translate function call
transFunCall :: String -> FunType -> [Expr] -> Trans (Expr, Type)
transFunCall name (FunType expectedTypes retType) args = do
    (args', inferredTypes) <- transExprs args
    let castArgs = [cast it et e | it <- inferredTypes |
            et <- expectedTypes | e <- args']
    return (Func name castArgs, retType)


-- translate user-defined function call
transUserCall :: String -> [Expr] -> Trans (Expr, Type)
transUserCall name args = do
    (args', inferredTypes) <- transExprs args
    env <- get
    case env # tenv # Map.lookup name of
        Nothing -> do
            -- function not seen before, translate it
            transUserFunction name args' inferredTypes
        Just f ->
            -- do the argument types match?
            if map snd (f # funArgs) == inferredTypes
                -- yes, just use the compiled function
                then return (Func name args', f # funType)
                -- no, have to compile a version with different arguments
                else transUserFunction name args' inferredTypes
                    

transUserFunction :: String -> [Expr] -> [Type] -> Trans (Expr, Type)
transUserFunction name args inferredTypes = do
    -- get function from uenv
    fun <- gets $ \env -> env # uenv # Map.lookup name
    case fun of
        Nothing -> error $ "function not found: " ++ name
        Just f@(Function _ _ argList body) -> 
            case length args == length argList of
                False -> error $ "wrong number of arguments in function call: " ++ name
                True -> do
                    -- specify argument types
                    let argList' = [(a, t) | a <- map fst argList | t <- inferredTypes]
            
                    -- translate function
                    fTrans <- transFunction (f {funArgs = argList'})
            
                    -- add it into tenv
                    tenv' <- gets $ \env -> env # tenv # Map.insert name fTrans
                    modify $ \env -> env {tenv = tenv'}
            
                    -- the return type will become the type of this expression
                    return (Func name args, fTrans # funType)


-- translate statement (potentially creating multiple statements)
transStmt :: Stmt -> Trans Stmt

transStmt (Group s) = do
    s' <- forM s transStmt
    return (Group s')

transStmt (Return e _) = do
    (e', t) <- transExpr e
    return (Return e' t)

transStmt (Assign x e) = do
    (e', inferredType) <- transExpr e
    env <- get
    case env # lenv # Map.lookup x of
        Nothing -> do
            -- add type declaration and insert inferred type into environment
            let lenv' = env # lenv # Map.insert x inferredType
            put (env {lenv = lenv'})
            return $ Group [TypeDecl x inferredType, Assign x e']
        Just previousType ->
            if inferredType == previousType
                -- OK, just continue
                then return (Assign x e')
                -- type does not match previously inferred, fail
                else error $ "type error in: " ++ printStmt (Assign x e)
        
transStmt (Branch e s1 s2) = do
    (e', t) <- transExpr e
    case t of
        Boolean -> do
            s1' <- transStmt s1
            s2' <- transStmt s2
            return (Branch e' s1' s2')
        _ -> error $ "expected boolean expression: " ++ printExpr e

transStmt (While e s) = do
    (e', t) <- transExpr e
    case t of
        Boolean -> do
            s' <- transStmt s
            return (While e' s')
        _ -> error $ "expected boolean expression: " ++ printExpr e

transStmt (For e u s) = do
    (e', t) <- transExpr e
    case t of
        Boolean -> do
            u' <- transStmt u
            s' <- transStmt s
            return (For e' u' s')
        _ -> error $ "expected boolean expression: " ++ printExpr e
        
transStmt Break = return Break

transStmt Continue = return Continue


transFunction :: Function -> Trans Function
transFunction f@(Function name retType args body) = do
    oldLocalEnv <- gets $ \env -> env # lenv
    
    -- add argument types into local environment
    modify $ \env -> env {lenv = Map.fromList args}
    
    -- translate body
    body' <- transStmt body
    
    -- remove redundant groups
    let body'' = removeGroups body'
    
    -- restore previous local environment
    modify $ \env -> env {lenv = oldLocalEnv}
    
    case retType of
        -- no desired return type, just derive one
        NoType -> return (Function name (getReturnType body'') args body'')
        -- otherwise, cast into the desired type
        t -> return (Function name t args (castReturnTypes t body''))
         

transProgram :: Program -> Trans Program

transProgram [f] = do
    -- this is the main function, run the translator on it
    f' <- transFunction f
    
    -- return the full list of translated functions
    env <- get
    let fs = env # tenv # Map.elems
    return (fs ++ [f'])
    
transProgram (f : fs) = do
    -- just add the function to uenv
    env <- get
    let uenv' = env # uenv # Map.insert (f # funName) f
    put (env {uenv = uenv'})
    transProgram fs
    

translate :: Program -> Program
translate p = evalTrans (transProgram p)



